/**
 * Dilithium多项式打包/解包 - NIST标准
 */

#include "packing.h"
#include "fips202.h"
#include <stdint.h>

/* 打包z多项式 (γ1 = 2^17) */
void polyz_pack(uint8_t *r, const Poly *a) {
    unsigned int i;
    uint32_t t[4];
    
    for(i = 0; i < DILITHIUM_N/4; ++i) {
        t[0] = (1 << (DILITHIUM_D - 1)) - a->coeffs[4*i+0];
        t[1] = (1 << (DILITHIUM_D - 1)) - a->coeffs[4*i+1];
        t[2] = (1 << (DILITHIUM_D - 1)) - a->coeffs[4*i+2];
        t[3] = (1 << (DILITHIUM_D - 1)) - a->coeffs[4*i+3];

        r[9*i+0] = t[0];
        r[9*i+1] = t[0] >> 8;
        r[9*i+2] = t[0] >> 16;
        r[9*i+2] |= t[1] << 2;
        r[9*i+3] = t[1] >> 6;
        r[9*i+4] = t[1] >> 14;
        r[9*i+4] |= t[2] << 4;
        r[9*i+5] = t[2] >> 4;
        r[9*i+6] = t[2] >> 12;
        r[9*i+6] |= t[3] << 6;
        r[9*i+7] = t[3] >> 2;
        r[9*i+8] = t[3] >> 10;
    }
}

/* 打包w1多项式 */
void polyw1_pack(uint8_t *r, const Poly *a) {
    unsigned int i;
    
    for(i = 0; i < DILITHIUM_N/2; ++i) {
        r[3*i+0] = a->coeffs[2*i+0];
        r[3*i+0] |= a->coeffs[2*i+1] << 6;
        r[3*i+1] = a->coeffs[2*i+1] >> 2;
        r[3*i+2] = a->coeffs[2*i+1] >> 10;
    }
}

/* 打包zVector */
void polyvecl_pack_z(uint8_t *r, const PolyVecL *v) {
    unsigned int i;
    for(i = 0; i < DILITHIUM_L; ++i)
        polyz_pack(r + i*DILITHIUM_POLYZ_PACKEDBYTES, &v->vec[i]);
}

/* 打包w1Vector */
void polyveck_pack_w1(uint8_t *r, const PolyVecK *v) {
    unsigned int i;
    for(i = 0; i < DILITHIUM_K; ++i)
        polyw1_pack(r + i*DILITHIUM_POLYW1_PACKEDBYTES, &v->vec[i]);
}

/* From种子Generate挑战多项式 */
void poly_challenge(Poly *c, const uint8_t seed[32]) {
    unsigned int i, b, pos;
    uint64_t signs;
    uint8_t buf[SHAKE256_RATE];
    keccak_state state;

    shake256_init(&state);
    shake256_absorb(&state, seed, 32);
    shake256_finalize(&state);
    shake256_squeeze(buf, sizeof(buf), &state);

    signs = 0;
    for(i = 0; i < 8; ++i)
        signs |= (uint64_t)buf[i] << 8*i;

    pos = 8;
    for(i = 0; i < DILITHIUM_N; ++i)
        c->coeffs[i] = 0;

    for(i = DILITHIUM_N-DILITHIUM_TAU; i < DILITHIUM_N; ++i) {
        do {
            if(pos >= SHAKE256_RATE) {
                shake256_squeeze(buf, sizeof(buf), &state);
                pos = 0;
            }
            b = buf[pos++];
        } while(b > i);

        c->coeffs[i] = c->coeffs[b];
        c->coeffs[b] = 1 - 2*(signs & 1);
        signs >>= 1;
    }
}
